# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
import torch.nn.functional as F

from mmseg.core import add_prefix
from mmseg.ops import resize
from .. import builder
from ..builder import SEGMENTORS
from .base import BaseSegmentor
from mmcv.runner import auto_fp16
from ..utils.CowMask import gen_cow_mask, gen_patch_mask, gen_mix_data
import numpy as np
import copy

from mmseg.utils.sam_tools import visualize_img_
from mmcv.image import tensor2imgs

@SEGMENTORS.register_module()
class EncoderDecoderUDAMix(BaseSegmentor):
    """Encoder Decoder segmentors for ST-DASegNet.

    EncoderDecoder_forDSFN typically consists of two backbone, two decode_head. Here, we do not
    apply auxiliary_head, neck to simplify the implementation.

    Args:
        backbone_s: backbone for source and target.
        decode_head: decode_head for source and target
    """

    def __init__(self,
                 backbone,
                 decode_head,
                 auxiliary_head=None,
                 discriminator_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None,
                 init_cfg=None,
                 full=False,
                 mix_style=None,
                 cross_EMA=None,
                 coco_mask=None):
        super(EncoderDecoderUDAMix, self).__init__(init_cfg)
        if pretrained is not None:
            assert backbone.get('pretrained') is None, \
                'both backbone and segmentor set pretrained weight'
            backbone.pretrained = pretrained
        self.backbone = builder.build_backbone(backbone)
        self._init_decode_head(decode_head)
        self._init_auxiliary_head(auxiliary_head)
        self._init_discriminator_head(discriminator_head)
        self.cross_EMA = cross_EMA
        if self.cross_EMA:
            self._init_teacher(cross_EMA)

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        self.full = full
        self.mix_style = mix_style
        self.coco_mask = coco_mask
        self.iteration = 0

        self._parse_train_cfg()

    def _init_teacher(self, cross_EMA):
        if not self.cross_EMA: return
        self.backbone_tea = copy.deepcopy(self.backbone)
        self.decode_head_tea = copy.deepcopy(self.decode_head)
        self.cross_EMA_alpha = 0.99
        return

    def _update_tea(self, iter):
        # alpha_tea = min(1 - 1 / (iter + 500), self.cross_EMA_alpha)

        # if iter < 5000:
        #     alpha_tea = (iter + 1) / 5000 * self.cross_EMA_alpha
        # else:
        #     alpha_tea = self.cross_EMA_alpha

        alpha_tea = self.cross_EMA_alpha

        """1. update backbone"""
        for tea_module, stu_module in zip(self.backbone_tea.parameters(), self.backbone.parameters()):

            # For scalar params slow
            if not tea_module.data.shape:
                tea_module.data = alpha_tea * tea_module.data + (1 - alpha_tea) * stu_module.data
                # For tensor params
            else:
                tea_module.data[:] = alpha_tea * tea_module.data[:] + (1 - alpha_tea) * stu_module.data[:]

            # # For scalar params fast
            # if not tea_module.data.shape:
            #     tea_module.data = alpha_tea * stu_module.data + (1 - alpha_tea) * tea_module.data
            #     # For tensor params
            # else:
            #     tea_module.data[:] = alpha_tea * stu_module.data[:] + (1 - alpha_tea) * tea_module.data[:]

        """2. update decode head"""
        # for tea_module, stu_module in zip(self.decode_head_tea.parameters(), self.decode_head.parameters()):
        #
        #     # For scalar params slow
        #     if not tea_module.data.shape:
        #         tea_module.data = alpha_tea * tea_module.data + (1 - alpha_tea) * stu_module.data
        #         # For tensor params
        #     else:
        #         tea_module.data[:] = alpha_tea * tea_module.data[:] + (1 - alpha_tea) * stu_module.data[:]
        #
        #     # # For scalar params fast
        #     # if not tea_module.data.shape:
        #     #     tea_module.data = alpha_tea * stu_module.data + (1 - alpha_tea) * tea_module.data
        #     #     # For tensor params
        #     # else:
        #     #     tea_module.data[:] = alpha_tea * stu_module.data[:] + (1 - alpha_tea) * tea_module.data[:]
        return alpha_tea

    def _get_tea_predict(self, img):
        f = self._forward_backbone(self.backbone_tea, img)
        # seg_logit_ori = self._forward_decode_head(self.decode_head_tea, f)
        seg_logit_ori = self._forward_decode_head(self.decode_head, f)

        seg_logit = resize(
            seg_logit_ori,
            size=img.shape[-2:],
            mode='bilinear',
            align_corners=self.align_corners,
            warning=False)

        psd_tea = seg_logit.max(1)[1].unsqueeze(1).float().detach()
        return psd_tea, seg_logit_ori

    def _get_self_predict(self, img):
        f = self._forward_backbone(self.backbone, img)
        seg_logit = self._forward_decode_head(self.decode_head, f)

        seg_logit = resize(
            seg_logit,
            size=img.shape[-2:],
            mode='bilinear',
            align_corners=self.align_corners,
            warning=False)

        psd = seg_logit.max(1)[1].unsqueeze(1).float().detach()
        return psd, seg_logit

    def _parse_train_cfg(self):
        """Parsing train config and set some attributes for training."""
        if self.train_cfg is None:
            self.train_cfg = dict()
        # control the work flow in train step
        self.disc_steps = self.train_cfg.get('disc_steps', 1)

        self.disc_init_steps = (0 if self.train_cfg is None else
                                self.train_cfg.get('disc_init_steps', 0))

    def _init_decode_head(self, decode_head):
        """Initialize ``decode_head``"""
        self.decode_head = builder.build_head(decode_head)
        self.align_corners = self.decode_head.align_corners
        self.num_classes = self.decode_head.num_classes

    def _init_auxiliary_head(self, auxiliary_head):
        """Initialize ``auxiliary_head``"""
        if auxiliary_head is not None:
            if isinstance(auxiliary_head, list):
                self.auxiliary_head = nn.ModuleList()
                for head_cfg in auxiliary_head:
                    self.auxiliary_head.append(builder.build_head(head_cfg))
            else:
                self.auxiliary_head = builder.build_head(auxiliary_head)

    def _init_discriminator_head(self, discriminator_head):
        """Initialize ``auxiliary_head``"""
        if discriminator_head is not None:
            if isinstance(discriminator_head, list):
                self.discriminator_head = nn.ModuleList()
                for head_cfg in discriminator_head:
                    self.discriminator_head.append(builder.build_head(head_cfg))
            else:
                self.discriminator_head = builder.build_head(discriminator_head)
        else:
            self.discriminator_head = None

    def extract_feat(self, img):
        """Extract features from images."""
        x = self.backbone(img)
        if self.with_neck:
            x = self.neck(x)
        return x

    def encode_decode(self, img, img_metas=None):
        """Encode images with backbone and decode into a semantic segmentation
        map of the same size as input."""
        # 1. encode features
        feature = self.extract_feat(img)
        # 2. decode output
        output = self._forward_decode_head(self.decode_head, feature)
        out = resize(
            input=output,
            size=img.shape[2:],
            mode='bilinear',
            align_corners=self.align_corners)
        return out

    def _decode_head_forward_test(self, x, img_metas):
        """Run forward function and calculate loss for decode head in
        inference."""
        seg_logits = self.decode_head.forward_test(x, img_metas, self.test_cfg)
        return seg_logits

    def forward_dummy(self, img):
        """Dummy forward function."""
        seg_logit = self.encode_decode(img, None)

        return seg_logit

    @staticmethod
    def _forward_backbone(backbone, img):
        F_b = backbone(img)
        return F_b

    @staticmethod
    def _forward_decode_head(decode_head, feature):
        Pred = decode_head(feature)
        return Pred

    def forward_train(self, batch_data, img_metas=None, **kwargs):
        pass
        """Forward function for training."""

    def _get_discriminator_loss(self, decode_head, pred, gt_domain, gt_weight=None):

        # print(f"pred shape : {pred.shape}\t gt_domain : {gt_domain.shape}")
        # raise  None

        gt_domain = F.interpolate(gt_domain.unsqueeze(1), size=pred.shape[-2:], mode='nearest').long()

        losses = dict()
        loss_seg = decode_head.losses(pred, gt_domain, gt_weight=gt_weight)
        losses.update(loss_seg)
        loss_seg, log_vars_seg = self._parse_losses(losses)
        return loss_seg, log_vars_seg

    def _get_segmentor_loss(self, decode_head, pred, gt_semantic_seg, gt_weight=None):
        losses = dict()
        loss_seg = decode_head.losses(pred, gt_semantic_seg, gt_weight=gt_weight)
        losses.update(loss_seg)
        loss_seg, log_vars_seg = self._parse_losses(losses)
        return loss_seg, log_vars_seg

    def _get_KD_loss(self, teacher, student, T=3):
        losses = dict()
        losses[f'loss'] = self.KL_loss(teacher, student, T)
        loss_KD, log_vars_KD = self._parse_losses(losses)
        return loss_KD, log_vars_KD

    def _get_mask(self, imgs):
        """
        self.mix_style is None or
        self.mix_style:
        { name:'cowmix',args:{xxx:xxx}
        }

        """
        assert self.mix_style is None or self.mix_style['name'] in ['cowmix', 'cowout', 'patchmix', 'patchout']
        if self.mix_style['name'] == 'cowmix' or self.mix_style['name'] == 'cowout':
            mixup_mask = (gen_cow_mask(imgs, self.mix_style['sigma']) > self.mix_style['threshold']).float()
        elif self.mix_style['name'] == 'patchmix' or 'patchout':
            mixup_mask = (gen_patch_mask(imgs, self.mix_style['block_size']) > self.mix_style['threshold']).float()
        else:
            raise ValueError(f'Invalid mix_style')

        return mixup_mask

    @auto_fp16()
    def train_step(self, data_batch, optimizer, **kwargs):
        """The iteration step during training.

        The whole process including back propagation and
        optimizer updating is also defined in this method, such as GAN.

        Args:
            data (dict): The output of dataloader.
            optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
                runner is passed to ``train_step()``. This argument is unused
                and reserved.

        Returns:
            dict: It should contain at least 3 keys: ``loss``, ``log_vars``,
                ``num_samples``.
                ``loss`` is a tensor for back propagation, which can be a
                weighted sum of multiple losses.
                ``log_vars`` contains all the variables to be sent to the
                logger.
                ``num_samples`` indicates the batch size (when the model is
                DDP, it means the batch size on each GPU), which is used for
                averaging the logs.
        """

        # dirty walkround for not providing running status
        if not hasattr(self, 'iteration'):
            self.iteration = 0
        curr_iter = self.iteration

        img_s = data_batch['img']
        gt_s = data_batch['gt_semantic_seg']
        img_t = data_batch['B_img']
        gt_t = data_batch['B_gt_semantic_seg']

        # print(img_s.shape,gt_s.shape)
        # print(img_t.shape,gt_t.shape)

        log_vars = dict()

        optimizer['backbone'].zero_grad()
        optimizer['decode_head'].zero_grad()

        # 1. forward source
        f_s = self._forward_backbone(self.backbone, img_s)
        out_s = self._forward_decode_head(self.decode_head, f_s)
        loss_s, log_vars_seg_s = self._get_segmentor_loss(self.decode_head, out_s, gt_s, None)
        log_vars['loss_s'] = log_vars_seg_s.pop('loss')
        # loss_s.backward(retain_graph=True)
        loss = loss_s

        # 2. forward aux decode head
        if self.with_auxiliary_head:
            out_s_aux = self._forward_decode_head(self.auxiliary_head, f_s)
            loss_s_aux, log_vars_seg_s_aux = self._get_segmentor_loss(self.auxiliary_head, out_s_aux, gt_s, None)
            log_vars['loss_s_aux'] = log_vars_seg_s_aux.pop('loss')
            # loss_s_aux.backward()
            loss += loss_s_aux

        if self.mix_style is not None or self.discriminator_head:

            if self.cross_EMA:
                psd_t, _ = self._get_tea_predict(img_t)
            else:
                psd_t, _ = self._get_self_predict(img_t)

            if self.coco_mask is not None and np.random.random() < 0.5:
                mixup_mask = data_batch['B_auto_mask']
            else:
                mixup_mask = self._get_mask(img_s)

            mixup_mask_copy = copy.deepcopy(mixup_mask)
            mixup_mask_copy[mixup_mask_copy == 255.] = 0.

            if np.random.random() < 0.5:
                mixup_mask_copy = 1 - mixup_mask_copy
            img1_mix, label1_mix, img2_mix, label2_mix = gen_mix_data(img_s, gt_s, img_t, psd_t, mode=self.mix_style['name'], mask=mixup_mask_copy)

            f1_mix = self._forward_backbone(self.backbone, img1_mix)
            f2_mix = self._forward_backbone(self.backbone, img2_mix)

        if self.mix_style is not None:
            out1_mix = self._forward_decode_head(self.decode_head, f1_mix)
            out2_mix = self._forward_decode_head(self.decode_head, f2_mix)

            loss_mix1, log_vars_seg_mix1 = self._get_segmentor_loss(self.decode_head, out1_mix, label1_mix, None)
            loss_mix2, log_vars_seg_mix2 = self._get_segmentor_loss(self.decode_head, out2_mix, label2_mix, None)

            log_vars['loss_mix'] = (log_vars_seg_mix1.pop('loss') + log_vars_seg_mix2.pop('loss')) / 2.
            loss += (loss_mix1 + loss_mix2) / 2.

        if self.cross_EMA:
            _, out1_tea = self._get_tea_predict(img1_mix)
            _, out2_tea = self._get_tea_predict(img2_mix)

            loss_kl1, log_vars_kl1 = self._get_KD_loss(out1_tea.detach(), out1_mix)
            loss_kl2, log_vars_kl2 = self._get_KD_loss(out2_tea.detach(), out2_mix)
            log_vars['loss_kl'] = (log_vars_kl1.pop('loss') + log_vars_kl2.pop('loss')) / 2.
            loss += (loss_kl1 + loss_kl2) / 2.

        if self.discriminator_head:
            out1_mix = self._forward_decode_head(self.discriminator_head, f1_mix)
            out2_mix = self._forward_decode_head(self.discriminator_head, f2_mix)
            loss_domain_mix1, log_vars_domain_mix1 = self._get_discriminator_loss(self.discriminator_head, out1_mix, mixup_mask, None)
            loss_domain_mix2, log_vars_domain_mix2 = self._get_discriminator_loss(self.discriminator_head, out2_mix, (1 - mixup_mask), None)

            log_vars['loss_domain'] = (log_vars_domain_mix1.pop('loss') + log_vars_domain_mix2.pop('loss')) / 2.
            loss += (loss_domain_mix1 + loss_domain_mix2) / 2.

        if self.full:
            f_t = self._forward_backbone(self.backbone, img_t)
            out_t = self._forward_decode_head(self.decode_head, f_t)
            loss_t, log_vars_seg_t = self._get_segmentor_loss(self.decode_head, out_t, gt_t, None)
            log_vars['loss_t'] = log_vars_seg_t.pop('loss')
            # loss_t.backward()
            loss += loss_t


        # PALETTE = [[255, 255, 255], [0, 0, 255], [0, 255, 255], [0, 255, 0],
        #            [255, 255, 0], [255, 0, 0]]

        # viz_img_s = tensor2imgs(img_s[:1,:,:,:],**data_batch['img_metas'][0]['img_norm_cfg'])[0]
        # viz_gt_s = gt_s[0,:,:,:].data.cpu().numpy().squeeze()
        # viz_img_t = tensor2imgs(img_t[:1,:,:,:],**data_batch['img_metas'][0]['img_norm_cfg'])[0]
        # viz_gt_t = gt_t[0,:,:,:].data.cpu().numpy().squeeze()
        # viz_psd_t = psd_t[0,:,:,:].data.cpu().numpy().squeeze()
        # viz_img1_mix=tensor2imgs(img1_mix[:1,:,:,:],**data_batch['img_metas'][0]['img_norm_cfg'])[0]
        # viz_label1_mix=label1_mix[0,:,:,:].data.cpu().numpy().squeeze()
        # viz_pred1_mix= resize(
        #     out1_mix,
        #     size=img_s.shape[-2:],
        #     mode='bilinear',
        #     align_corners=self.align_corners,
        #     warning=False).max(1)[1].unsqueeze(1).float().detach().data.cpu().numpy().squeeze()[0]
        # viz_mixup_mask=mixup_mask[0].data.cpu().numpy().squeeze()
        # import os
        # visualize_img_(viz_img_s[:,:,::-1],save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_img_s.png'))
        # visualize_img_(viz_img_t[:,:,::-1],save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_img_t.png'))
        # visualize_img_(viz_img1_mix[:,:,::-1],save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_img1_mix.png'))
        #
        # visualize_img_(viz_gt_s,palette=PALETTE,save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_gt_s.png'))
        # visualize_img_(viz_gt_t,palette=PALETTE,save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_gt_t.png'))
        # visualize_img_(viz_psd_t,palette=PALETTE,save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_psd_t.png'))
        # visualize_img_(viz_label1_mix,palette=PALETTE,save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_label1_mix.png'))
        # visualize_img_(viz_pred1_mix,palette=PALETTE,save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_pred1_mix.png'))
        #
        # visualize_img_(viz_mixup_mask,palette=[[0,0,0],[255,255,255]],save_path=os.path.join('./viz_dir/patch_cocomask_viz','viz_mixup_mask.png'))


        loss.backward()
        optimizer['backbone'].step()
        optimizer['decode_head'].step()

        if self.cross_EMA:
            alpha_tea = self._update_tea(self.iteration)
            log_vars['alpha_tea'] = alpha_tea

        if hasattr(self, 'iteration'):
            self.iteration += 1
        outputs = dict(
            loss=loss,
            log_vars=log_vars,
            num_samples=len(data_batch['img_metas']))

        return outputs

    # TODO refactor
    def slide_inference(self, img, img_meta, rescale):
        """Inference by sliding-window with overlap.

        If h_crop > h_img or w_crop > w_img, the small patch will be used to
        decode without padding.
        """

        h_stride, w_stride = self.test_cfg.stride
        h_crop, w_crop = self.test_cfg.crop_size
        batch_size, _, h_img, w_img = img.size()
        num_classes = self.num_classes
        h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
        w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
        preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
        count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
        for h_idx in range(h_grids):
            for w_idx in range(w_grids):
                y1 = h_idx * h_stride
                x1 = w_idx * w_stride
                y2 = min(y1 + h_crop, h_img)
                x2 = min(x1 + w_crop, w_img)
                y1 = max(y2 - h_crop, 0)
                x1 = max(x2 - w_crop, 0)
                crop_img = img[:, :, y1:y2, x1:x2]
                crop_seg_logit = self.encode_decode(crop_img, img_meta)
                preds += F.pad(crop_seg_logit,
                               (int(x1), int(preds.shape[3] - x2), int(y1),
                                int(preds.shape[2] - y2)))

                count_mat[:, :, y1:y2, x1:x2] += 1
        assert (count_mat == 0).sum() == 0
        if torch.onnx.is_in_onnx_export():
            # cast count_mat to constant while exporting to ONNX
            count_mat = torch.from_numpy(
                count_mat.cpu().detach().numpy()).to(device=img.device)
        preds = preds / count_mat
        if rescale:
            preds = resize(
                preds,
                size=img_meta[0]['ori_shape'][:2],
                mode='bilinear',
                align_corners=self.align_corners,
                warning=False)
        return preds

    def whole_inference(self, img, img_meta, rescale):
        """Inference with full image."""

        seg_logit = self.encode_decode(img, img_meta)
        if rescale:
            # support dynamic shape for onnx
            if torch.onnx.is_in_onnx_export():
                size = img.shape[2:]
            else:
                size = img_meta[0]['ori_shape'][:2]
            seg_logit = resize(
                seg_logit,
                size=size,
                mode='bilinear',
                align_corners=self.align_corners,
                warning=False)

        return seg_logit

    def inference(self, img, img_meta, rescale):
        """Inference with slide/whole style.

        Args:
            img (Tensor): The input image of shape (N, 3, H, W).
            img_meta (dict): Image info dict where each dict has: 'img_shape',
                'scale_factor', 'flip', and may also contain
                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
                For details on the values of these keys see
                `mmseg/datasets/pipelines/formatting.py:Collect`.
            rescale (bool): Whether rescale back to original shape.

        Returns:
            Tensor: The output segmentation map.
        """

        assert self.test_cfg.mode in ['slide', 'whole']
        ori_shape = img_meta[0]['ori_shape']
        assert all(_['ori_shape'] == ori_shape for _ in img_meta)
        if self.test_cfg.mode == 'slide':
            seg_logit = self.slide_inference(img, img_meta, rescale)
        else:
            seg_logit = self.whole_inference(img, img_meta, rescale)
        output = F.softmax(seg_logit, dim=1)
        flip = img_meta[0]['flip']
        if flip:
            flip_direction = img_meta[0]['flip_direction']
            assert flip_direction in ['horizontal', 'vertical']
            if flip_direction == 'horizontal':
                output = output.flip(dims=(3,))
            elif flip_direction == 'vertical':
                output = output.flip(dims=(2,))

        return output

    def simple_test(self, img, img_meta, rescale=True):
        """Simple test with single image."""
        seg_logit = self.inference(img, img_meta, rescale)
        seg_pred = seg_logit.argmax(dim=1)
        if torch.onnx.is_in_onnx_export():
            # our inference backend only support 4D output
            seg_pred = seg_pred.unsqueeze(0)
            return seg_pred
        seg_pred = seg_pred.cpu().numpy()
        # unravel batch dim
        seg_pred = list(seg_pred)
        return seg_pred

    def aug_test(self, imgs, img_metas, rescale=True):
        """Test with augmentations.

        Only rescale=True is supported.
        """
        # aug_test rescale all imgs back to ori_shape for now
        assert rescale
        # to save memory, we get augmented seg logit inplace
        seg_logit = self.inference(imgs[0], img_metas[0], rescale)
        for i in range(1, len(imgs)):
            cur_seg_logit = self.inference(imgs[i], img_metas[i], rescale)
            seg_logit += cur_seg_logit
        seg_logit /= len(imgs)
        seg_pred = seg_logit.argmax(dim=1)
        seg_pred = seg_pred.cpu().numpy()
        # unravel batch dim
        seg_pred = list(seg_pred)
        return seg_pred

    ## added by LYU: 2022/05/12
    def MSE_loss(self, teacher, student):
        MSE_loss = nn.MSELoss()
        t = self.sw_softmax(teacher)
        s = self.sw_softmax(student)
        KD_loss = MSE_loss(s, t)
        return KD_loss

    @staticmethod
    def set_requires_grad(nets, requires_grad=False):
        """Set requires_grad for all the networks.

        Args:
            nets (nn.Module | list[nn.Module]): A list of networks or a single
                network.
            requires_grad (bool): Whether the networks require gradients or not
        """
        if not isinstance(nets, list):
            nets = [nets]
        for net in nets:
            if net is not None:
                for param in net.parameters():
                    param.requires_grad = requires_grad

    @staticmethod
    def sw_softmax(pred):
        N, C, H, W = pred.shape
        pred_sh = torch.reshape(pred, (N, C, H * W))
        pred_sh = F.softmax(pred_sh, dim=2)
        pred_out = torch.reshape(pred_sh, (N, C, H, W))
        return pred_out

    ## added by LYU: 2022/05/11
    @staticmethod
    def KL_loss(teacher, student, T=5):
        KL_loss = nn.KLDivLoss(reduction='mean')(F.log_softmax(student / T, dim=1),
                                                 F.softmax(teacher / T, dim=1)) * (T * T)
        return KL_loss
